#!/usr/bin/env python3
"""
Add code_content attribute to graph nodes by copying from code attribute
"""

import pickle
import os
from collections import defaultdict

# Input and output file paths
input_pkl = "/data_ext/ref_code/LocAgent/index_data_strict/pp_data/graph_index_v2.3/business-payout.pkl"
output_pkl = "/data_ext/ref_code/LocAgent/index_data_strict/pp_data/graph_index_v2.3/business-payout-with-code-content.pkl"

print("Loading graph from pickle file...")
with open(input_pkl, "rb") as f:
    G = pickle.load(f)

print(f"\n=== Original Graph Statistics ===")
print(f"Total nodes: {G.number_of_nodes()}")
print(f"Total edges: {G.number_of_edges()}")

# Statistics
nodes_modified = 0
nodes_with_code = 0
nodes_without_code = 0
node_types_modified = defaultdict(int)

print(f"\n=== Adding code_content attribute ===")

# Iterate through all nodes and add code_content
for node_id in G.nodes():
    node_data = G.nodes[node_id]

    # If node has 'code' attribute, copy it to 'code_content'
    if "code" in node_data:
        # Copy code to code_content
        G.nodes[node_id]["code_content"] = node_data["code"]
        nodes_modified += 1
        nodes_with_code += 1

        # Track which types were modified
        if "type" in node_data:
            node_types_modified[node_data["type"]] += 1

        # Log progress every 500 nodes
        if nodes_modified % 500 == 0:
            print(f"  Processed {nodes_modified} nodes...")
    else:
        nodes_without_code += 1

print(f"\n=== Modification Summary ===")
print(f"Nodes with code (modified): {nodes_with_code}")
print(f"Nodes without code: {nodes_without_code}")
print(f"Total nodes processed: {nodes_with_code + nodes_without_code}")

print(f"\n=== Modified Node Types ===")
for node_type, count in sorted(node_types_modified.items()):
    print(f"{node_type}: {count} nodes")

# Save the modified graph
print(f"\n=== Saving modified graph ===")
with open(output_pkl, "wb") as f:
    pickle.dump(G, f)
print(f"Saved to: {output_pkl}")

# Verify the modification by checking a few samples
print(f"\n=== Verification ===")
sample_count = 0
for node_id in G.nodes():
    node_data = G.nodes[node_id]
    if "code_content" in node_data and "code" in node_data:
        if sample_count == 0:
            print(f"\nSample node: {node_id}")
            print(f"  Has 'code': {('code' in node_data)}")
            print(f"  Has 'code_content': {('code_content' in node_data)}")
            print(
                f"  Content identical: {node_data.get('code') == node_data.get('code_content')}"
            )

            if node_data.get("code_content"):
                lines = node_data["code_content"].split("\n")
                print(
                    f"  Code length: {len(lines)} lines, {len(node_data['code_content'])} chars"
                )
        sample_count += 1
        if sample_count >= 5:
            break

print(f"\nVerified {sample_count} nodes have both 'code' and 'code_content' attributes")
print(f"\n✅ Graph modification completed successfully!")
print(f"Original file: {input_pkl}")
print(f"Modified file: {output_pkl}")
print(f"File size: {os.path.getsize(output_pkl) / (1024*1024):.2f} MB")
